// Copyright DustedPixels.com 2008. All rights reserved.

package com.dustedpixels.asm.utils;

import java.util.HashMap;
import java.util.HashSet;
import java.util.Map;
import java.util.Set;

import org.objectweb.asm.MethodAdapter;
import org.objectweb.asm.MethodVisitor;
import org.objectweb.asm.Opcodes;
import org.objectweb.asm.Type;
import org.objectweb.asm.commons.AdviceAdapter;
import org.objectweb.asm.tree.AbstractInsnNode;
import org.objectweb.asm.tree.FieldInsnNode;
import org.objectweb.asm.tree.InsnNode;
import org.objectweb.asm.tree.MethodNode;
import org.objectweb.asm.tree.VarInsnNode;
import org.objectweb.asm.tree.analysis.Analyzer;
import org.objectweb.asm.tree.analysis.AnalyzerException;
import org.objectweb.asm.tree.analysis.Frame;
import org.objectweb.asm.tree.analysis.Interpreter;
import org.objectweb.asm.tree.analysis.SourceInterpreter;
import org.objectweb.asm.tree.analysis.SourceValue;

import com.dustedpixels.asm.OpcodeUtils;
import com.dustedpixels.asm.analysis.Analyzers;
import com.dustedpixels.asm.analysis.FlowInterpreter;
import com.dustedpixels.asm.analysis.FlowValue;
import com.dustedpixels.asm.analysis.FrameUtils;
import com.dustedpixels.asm.analysis.TargetInterpreter;
import com.dustedpixels.asm.analysis.TargetValue;


/**
 * @author micapolos@gmail.com (Michal Pociecha-Los)
 */
public final class FieldLocalizer extends MethodAdapter {
  public static MethodVisitor get(MethodVisitor out, String className, String methodDesc) {
    return new FieldLocalizer(out, className, methodDesc);
  }
  
  private final MethodVisitor out;
  private final String className;

  FieldLocalizer(MethodVisitor out, 
      String className,
      String methodDesc) {
    
    super(new MethodNode(0, "foo", methodDesc, null, new String[0]));
    
    this.out = out;
    this.className = className;
  }

  @Override
  public void visitEnd() {
    super.visitEnd();
    
    MethodNode method = (MethodNode) this.mv;
    
    try {
      Interpreter flowInterpreter = new FlowInterpreter();
      Interpreter sourceInterpreter = new SourceInterpreter();
      Interpreter targetInterpreter = new TargetInterpreter();
      
      Analyzer flowAnalyzer = Analyzers.newUsing(flowInterpreter);
      Analyzer sourceAnalyzer = Analyzers.newUsing(sourceInterpreter);
      Analyzer targetAnalyzer = Analyzers.newUsing(targetInterpreter);
      
      System.out.println("Flow analysis");
      Frame[] flowFrames = flowAnalyzer.analyze(className, method);
      System.out.println("Source analysis");
      Frame[] sourceFrames = sourceAnalyzer.analyze(className, method);
      System.out.println("Target analysis");
      Frame[] targetFrames = targetAnalyzer.analyze(className, method);
      System.out.println("Analysis ended");
      
      AbstractInsnNode[] insns = method.instructions.toArray();

      // Collect all GETFIELD and PUTFIELD instructions, which 
      final Set<FieldInsnNode> fieldInsns = new HashSet<FieldInsnNode>();
      final Map<String, String> fieldName2DescMap = new HashMap<String, String>();
      final Set<String> uncertainFieldNames = new HashSet<String>();
      
      for (int i = 0; i < insns.length; i++) {
        AbstractInsnNode insn = insns[i];
        if (insn instanceof FieldInsnNode) {
          FieldInsnNode fieldInsn = (FieldInsnNode) insn;
          if (insn.getOpcode() == Opcodes.GETFIELD || insn.getOpcode() == Opcodes.PUTFIELD) {
            int stackIndex = (insn.getOpcode() == Opcodes.GETFIELD) ? 0 : 1;
            Frame flowFrame = flowFrames[i];
            FlowValue flowValue = (FlowValue) FrameUtils.top(flowFrame, stackIndex);
            FlowValue thisValue = (FlowValue) flowFrames[0].getLocal(0);
            if (flowValue == thisValue) {
              //System.out.println("Field access at: " + i);
              fieldInsns.add(fieldInsn);
              fieldName2DescMap.put(fieldInsn.name, fieldInsn.desc);
            } else if (flowValue.covers(thisValue)) {
              //System.out.println("Uncertain field access at: " + i);
              uncertainFieldNames.add(fieldInsn.name);
            }
          }
        }
      }
      
      // Replace field instructions with load/store.
      final Set<AbstractInsnNode> insnsToRemove = new HashSet<AbstractInsnNode>();
      final Set<FieldInsnNode> fieldInsnsWithPop = new HashSet<FieldInsnNode>();
      final Set<String> fieldNames = new HashSet<String>();
      final Set<String> getFieldNames = new HashSet<String>();
      final Set<String> putFieldNames = new HashSet<String>();
      
      for (FieldInsnNode fieldInsn : fieldInsns) {
        if (!uncertainFieldNames.contains(fieldInsn.name)) {
          fieldNames.add(fieldInsn.name);
          if (fieldInsn.getOpcode() == Opcodes.GETFIELD) {
            getFieldNames.add(fieldInsn.name);
          } else {
            putFieldNames.add(fieldInsn.name);
          }
          System.out.println("Processing field " + fieldInsn.name);
          int fieldInsnIndex = method.instructions.indexOf(fieldInsn);
          System.out.println("...at " + fieldInsnIndex);
          int stackIndex = (fieldInsn.getOpcode() == Opcodes.GETFIELD) ? 0 : 1;
          Frame sourceFrame = sourceFrames[fieldInsnIndex];
          SourceValue sourceValue = (SourceValue) FrameUtils.top(sourceFrame, stackIndex);
          AbstractInsnNode sourceInsnToRemove = null;
          System.out.println(" Number of sources: " + sourceValue.insns.size());
          if (sourceValue.insns.size() == 1) {
            AbstractInsnNode sourceInsn = (AbstractInsnNode) sourceValue.insns.iterator().next();
            System.out.println(" Source instruction found at: " + method.instructions.indexOf(sourceInsn));
            if (sourceInsn.getOpcode() == Opcodes.ALOAD ||
                sourceInsn.getOpcode() == Opcodes.DUP) {
              int sourceInsnIndex = method.instructions.indexOf(sourceInsn);
              TargetValue targetValue = (TargetValue) FrameUtils.top(targetFrames[sourceInsnIndex + 1]);
              System.out.println(" Number of targets: " + targetValue.getTargets().size());
              int consumingTargets = 0;
              for (AbstractInsnNode targetInsn : targetValue.getTargets()) {
                switch (targetInsn.getOpcode()) {
                  case Opcodes.DUP:
                  case Opcodes.DUP2:
                  case Opcodes.DUP2_X1:
                  case Opcodes.DUP2_X2:
                  case Opcodes.DUP_X1:
                  case Opcodes.DUP_X2:
                    break;
                  default:
                    consumingTargets++;
                }
              }
              System.out.println(" Number of consuming targets: " + consumingTargets);
              if (consumingTargets == 1) {
                sourceInsnToRemove = sourceInsn;
              }
            }
          }
  
          if (sourceInsnToRemove == null) {
            System.out.println("POP case");
            fieldInsnsWithPop.add(fieldInsn);
          } else {
            System.out.println("ALOAD/DUP case");
            if (insnsToRemove.contains(sourceInsnToRemove)) {
              throw new IllegalStateException();
            }
            insnsToRemove.add(sourceInsnToRemove);
          }
        }
      }
      
      // Generate locals for fields
      int nextLocal = method.maxLocals;
      final Map<String, Integer> fieldName2LocalMap = new HashMap<String, Integer>();
      for (String fieldName : fieldNames) {
        fieldName2LocalMap.put(fieldName, nextLocal);
        String fieldDesc = fieldName2DescMap.get(fieldName);
        Type fieldType = Type.getType(fieldDesc);
        int fieldSize = fieldType.getSize();
        nextLocal += fieldSize;
      }
      method.maxLocals = nextLocal;

      // Remove instructions.
      for (AbstractInsnNode insnToRemove : insnsToRemove) {
        method.instructions.remove(insnToRemove);
      }
      
      for (FieldInsnNode fieldInsn : fieldInsns) {
        VarInsnNode varInsn;
        if (fieldInsn.getOpcode() == Opcodes.GETFIELD) {
          varInsn = new VarInsnNode(
              OpcodeUtils.getLoadOpcode(Type.getType(fieldInsn.desc)), 
              fieldName2LocalMap.get(fieldInsn.name));
        } else {
          varInsn = new VarInsnNode(
              OpcodeUtils.getStoreOpcode(Type.getType(fieldInsn.desc)),
              fieldName2LocalMap.get(fieldInsn.name));
        }

        if (fieldInsnsWithPop.contains(fieldInsn)) {
          System.out.println(" Inserting POP");
          InsnNode popInsn = new InsnNode(Opcodes.POP);
          method.instructions.insertBefore(fieldInsn, popInsn);
        }
        method.instructions.insertBefore(fieldInsn, varInsn);
        method.instructions.remove(fieldInsn);
      }
      
      // Wrap method with field->local and local->field transfer
      MethodVisitor wrapper = new AdviceAdapter(out, method.access, method.name, method.desc) {
        @Override protected void onMethodEnter() {
          for (String fieldName : getFieldNames) {
            String fieldDesc = fieldName2DescMap.get(fieldName);
            super.visitVarInsn(Opcodes.ALOAD, 0);
            super.visitFieldInsn(Opcodes.GETFIELD, 
                className,
                fieldName, 
                fieldName2DescMap.get(fieldName));
            super.visitVarInsn(
                OpcodeUtils.getStoreOpcode(Type.getType(fieldDesc)),
                fieldName2LocalMap.get(fieldName));
          }
        }

        @Override
        protected void onMethodExit(int opcode) {
          for (String fieldName : putFieldNames) {
            String fieldDesc = fieldName2DescMap.get(fieldName);
            super.visitVarInsn(Opcodes.ALOAD, 0);
            super.visitVarInsn(
                OpcodeUtils.getLoadOpcode(Type.getType(fieldDesc)),
                fieldName2LocalMap.get(fieldName));
            super.visitFieldInsn(Opcodes.PUTFIELD, 
                className,
                fieldName, 
                fieldName2DescMap.get(fieldName));
          }
        }
      };
      
      method.instructions.resetLabels();
      
      method.accept(wrapper);
    } catch (AnalyzerException e) {
      throw new RuntimeException(e);
    }
  }
}
